#include <cstdlib>
#include <cstring>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <iterator>
#include <memory>

#include "protonos8.hpp"
#include "bitsetcount.hpp"

namespace protonos8 {

using boost::asio::deadline_timer;
using boost::asio::ip::udp;

client::client( const char* host, const char* port,
	boost::posix_time::time_duration timeout)
	:
	socket_(io_service_, udp::endpoint(udp::v4(), 0)),
	resolver_(io_service_),
	query_(udp::v4(), host, port),
	iterator_(resolver_.resolve(query_)),
	deadline_(io_service_),
	timeout_(timeout),
	fid_counters(boost::extents[256][8])
{
	std::fill(fid_counters.data(), fid_counters.data() + fid_counters.num_elements(), 0);
	deadline_.expires_at(boost::posix_time::pos_infin);
	check_deadline();
}

ack_result
client::info(uint32_t& data)
{
	ack_result result = no_answer;
	req_hdr r;
	pb_info_ack a;
	boost::system::error_code ec;

	r.fid = fid_info;
	r.addr.piggyback_location = 0;
	r.addr.counter = fid_counters[fid_info][0]++;

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	receive(boost::asio::buffer(&a, sizeof(a)), timeout_, ec);

	if (!ec) {
		if((a.hdr.fid == r.fid) &&
			(a.hdr.addr.piggyback_location == r.addr.piggyback_location) &&
			(a.hdr.addr.counter == r.addr.counter))
		{
			data = a.data;
			result = static_cast<ack_result>(a.hdr.result);
		} else
			result = wrong_ack;
	}


	return result;
}	

ack_result
client::load_blob_send_start_command(std::size_t& n)
{
	std::size_t retries = 3;
	ack_result result = no_answer;
	pb_load_blob_req_hdr r;
	pb_load_blob_ack a;
	boost::system::error_code ec;

	n = 0;
	r.fid = fid_another_load_blob;
	r.nested_fid = nested_fid_start;
	r.args = nested_fid_start_magic;

	do {
		socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

		receive(boost::asio::buffer(&a, sizeof(a)), boost::posix_time::seconds(3), ec);

		if (!ec) {
			if((a.fid == r.fid) &&
				(a.nested_fid == r.nested_fid))
			{
				n = a.next_line;
				result = static_cast<ack_result>(a.result);
			} else
				result = wrong_ack;
		}

	} while (result == no_answer && --retries);

	return result;
}

ack_result
client::load_blob_send_data(std::size_t& n, std::string& s)
{
	std::size_t retries = 3;
	ack_result result = no_answer;
	pb_load_blob_ack a;
	boost::system::error_code ec;

	if (s.size() > hex_max_size) s.resize(hex_max_size);

	std::size_t request_size = sizeof(pb_load_blob_req_hdr) + s.size();
	uint8_t* request = new uint8_t[request_size];
	pb_load_blob_req_hdr* r = reinterpret_cast<pb_load_blob_req_hdr*>(request);
	char* d = reinterpret_cast<char*>(&request[sizeof(pb_load_blob_req_hdr)]);

	r->fid = fid_another_load_blob;
	r->nested_fid = nested_fid_data;
	r->args = n;
	s.copy(d, s.size());

	do {
		socket_.send_to(boost::asio::buffer(request, request_size), *iterator_);

		receive(boost::asio::buffer(&a, sizeof(a)), boost::posix_time::seconds(3), ec);

		if (!ec) {
			if((a.fid == r->fid) &&
				(a.nested_fid == r->nested_fid))
			{
				n = a.next_line;
				result = static_cast<ack_result>(a.result);
			} else
				result = wrong_ack;
		}

	} while (result == no_answer && --retries);

	delete[] request;
	return result;
}

ack_result
client::load_blob_send_write_command(void)
{
	std::size_t retries = 3;
	ack_result result = no_answer;
	pb_load_blob_req_hdr r;
	pb_load_blob_ack a;
	boost::system::error_code ec;

	r.fid = fid_another_load_blob;
	r.nested_fid = nested_fid_write;
	r.args = nested_fid_write_magic;

	do {
		socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

		receive(boost::asio::buffer(&a, sizeof(a)), boost::posix_time::seconds(5), ec);

		if (!ec) {
			if((a.fid == r.fid) &&
				(a.nested_fid == r.nested_fid))
			{
				result = static_cast<ack_result>(a.result);
			} else
				result = wrong_ack;
		}

	} while (result == no_answer && --retries);

	return result;
}

ack_result
client::load_blob(char* fn)
{
	ack_result result = no_answer;
	std::vector<std::string> tokens(
		(std::istream_iterator<std::string>(
			*(std::auto_ptr<std::ifstream>(
				new std::ifstream(fn)
			)).get()
		)),
		std::istream_iterator<std::string>()
	);

	std::size_t next_line = 0;

	if ((result = load_blob_send_start_command(next_line)) != success)
		return result;

	if (next_line == 0xffff) {
		std::cout << "'Start' command is forbidden (wrong magic or flash is write protected)!\n";
		return wrong_ack;
	}

	while (next_line < tokens.size()) {
		result = load_blob_send_data(next_line, tokens[next_line]);
		if ((result != success) &&
				(result != out_of_order)) {
			std::cout << "'Data' command returned error " << result << std::endl;
			return result;
		}
		if (!(next_line % 16))
			std::cout << "[" << next_line << "]\r";
	}

	return load_blob_send_write_command();
}

ack_result
client::write_mmc01(uint8_t* out)
{
	ack_result result = no_answer;
	pb_write_mmc01_req r;
	ack_hdr a;
	boost::system::error_code ec;

	r.hdr.fid = fid_write_mmc01;
	r.hdr.addr.piggyback_location = 0;
	r.hdr.addr.counter = fid_counters[fid_write_mmc01][0]++;
	std::copy(out, out+8, &r.out[0]);

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	receive(boost::asio::buffer(&a, sizeof(a)), timeout_, ec);

	if (!ec) {
		if((a.fid == r.hdr.fid) &&
			(a.addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a.addr.counter == r.hdr.addr.counter))
			result = static_cast<ack_result>(a.result);
		else
			result = wrong_ack;
	}

	return result;
}

ack_result
client::raw_rw(function_id fid, unsigned pb_loc, uint32_t addr, uint32_t& value)
{
	ack_result result = no_answer;
	pb_raw_rw_req r;
	pb_raw_rw_ack a;
	boost::system::error_code ec;

	r.hdr.fid = fid;
	r.hdr.addr.piggyback_location = pb_loc;
	r.hdr.addr.counter = fid_counters[fid][pb_loc]++;
	r.addr = addr;	
	r.data = value;

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	receive(boost::asio::buffer(&a, sizeof(a)), timeout_, ec);

	if (!ec) {
		if((a.hdr.fid == r.hdr.fid) &&
			(a.hdr.addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a.hdr.addr.counter == r.hdr.addr.counter))
		{
			if (a.hdr.result == success)
				value = a.data;

			result = static_cast<ack_result>(a.hdr.result);
		} else
			result = wrong_ack;
	}

	return result;
}

ack_result
client::rw_mmc(function_id fid,
	uint8_t out_mask, uint16_t* dout,
	uint32_t inp_mask, uint16_t* dinp)
{
	ack_result result = no_answer;
	pb_mmc_rw_req r;
	boost::system::error_code ec;

	std::size_t len = bitsetcount(inp_mask);
	std::size_t ack_size = sizeof(ack_hdr) + len*sizeof(uint16_t);

	r.hdr.fid = fid;
	r.hdr.addr.piggyback_location = 0;
	r.hdr.addr.counter = fid_counters[fid][0]++;
	r.out_mask = out_mask;	
	r.inp_mask = inp_mask;
	std::copy(dout, dout+8, &r.dout[0]);

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	uint8_t* ack = new uint8_t[ack_size];
	ack_hdr* a = reinterpret_cast<ack_hdr*>(ack);
	uint16_t* d = reinterpret_cast<uint16_t*>(&ack[sizeof(ack_hdr)]);

	receive(boost::asio::buffer(ack, ack_size), timeout_, ec);

	if (!ec) {
		if((a->fid == r.hdr.fid) &&
			(a->addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a->addr.counter == r.hdr.addr.counter))
		{
			if (a->result == success)
				std::copy(d, d+len, dinp);
			result = static_cast<ack_result>(a->result);
		} else
			result = wrong_ack;
	}

	delete[] ack;
	return result;
}

ack_result
client::read_counters(function_id fid, unsigned pb_loc, uint8_t inp_chan, uint32_t* chn)
{
	ack_result result = no_answer;
	pb_counters_req r;
	boost::system::error_code ec;

	std::size_t len = (inp_chan == 0xff) ? 4 : 1;
	std::size_t ack_size = sizeof(ack_hdr) + len*sizeof(uint32_t);

	r.hdr.fid = fid;
	r.hdr.addr.piggyback_location = pb_loc;
	r.hdr.addr.counter = fid_counters[fid][pb_loc]++;
	r.inp_chan = inp_chan;

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	uint8_t* ack = new uint8_t[ack_size];
	ack_hdr* a = reinterpret_cast<ack_hdr*>(ack);
	uint32_t* d = reinterpret_cast<uint32_t*>(&ack[sizeof(ack_hdr)]);

	receive(boost::asio::buffer(ack, ack_size), timeout_, ec);

	if (!ec) {
		if((a->fid == r.hdr.fid) &&
			(a->addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a->addr.counter == r.hdr.addr.counter))
		{
			if (a->result == success)
				std::copy(d, d+len, chn);
			result = static_cast<ack_result>(a->result);
		} else
			result = wrong_ack;
	}

	delete[] ack;
	return result;
}

ack_result
client::read_adc(function_id fid, unsigned pb_loc, uint16_t inp_mask, uint16_t* dinp)
{
	ack_result result = no_answer;
	pb_adc_req r;
	boost::system::error_code ec;

	std::size_t len = bitsetcount(inp_mask & 0xffff);
	std::size_t ack_size = sizeof(ack_hdr) + len*sizeof(uint16_t);

	r.hdr.fid = fid;
	r.hdr.addr.piggyback_location = pb_loc;
	r.hdr.addr.counter = fid_counters[fid][pb_loc]++;
	r.inp_mask = inp_mask;

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	uint8_t* ack = new uint8_t[ack_size];
	ack_hdr* a = reinterpret_cast<ack_hdr*>(ack);
	uint16_t* d = reinterpret_cast<uint16_t*>(&ack[sizeof(ack_hdr)]);

	receive(boost::asio::buffer(ack, ack_size), timeout_, ec);

	if (!ec) {
		if((a->fid == r.hdr.fid) &&
			(a->addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a->addr.counter == r.hdr.addr.counter))
		{
			if (a->result == success)
				std::copy(d, d+len, dinp);
			result = static_cast<ack_result>(a->result);
		} else
			result = wrong_ack;
	}

	delete[] ack;
	return result;
}

ack_result
client::write_dac(function_id fid, unsigned pb_loc, uint8_t out_mask, uint16_t* dout)
{
	ack_result result = no_answer;
	pb_dac_req r;
	ack_hdr a;
	boost::system::error_code ec;

	r.hdr.fid = fid;
	r.hdr.addr.piggyback_location = pb_loc;
	r.hdr.addr.counter = fid_counters[fid][pb_loc]++;
	r.out_mask = out_mask;
	std::copy(dout, dout+4, &r.dout[0]);

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	receive(boost::asio::buffer(&a, sizeof(a)), timeout_, ec);

	if (!ec) {
		if((a.fid == r.hdr.fid) &&
			(a.addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a.addr.counter == r.hdr.addr.counter))
		{
			result =  static_cast<ack_result>(a.result);
		} else
			result = wrong_ack;
	}

	return result;
}

ack_result
client::init_location(function_id fid, unsigned pb_loc, uint8_t ini_set)
{
	ack_result result = no_answer;
	pb_init_location_req r;
	ack_hdr a;
	boost::system::error_code ec;

	r.hdr.fid = fid;
	r.hdr.addr.piggyback_location = pb_loc;
	r.hdr.addr.counter = fid_counters[fid][pb_loc]++;
	r.ini_set = ini_set;

	socket_.send_to(boost::asio::buffer(&r, sizeof(r)), *iterator_);

	receive(boost::asio::buffer(&a, sizeof(a)), timeout_, ec);

	if (!ec) {
		if((a.fid == r.hdr.fid) &&
			(a.addr.piggyback_location == r.hdr.addr.piggyback_location) &&
			(a.addr.counter == r.hdr.addr.counter))
		{
			result = static_cast<ack_result>(a.result);
		} else
			result = wrong_ack;
	}

	return result;

}

void
client::check_deadline()
{
	if (deadline_.expires_at() <= deadline_timer::traits_type::now()) {
		socket_.cancel();
		deadline_.expires_at(boost::posix_time::pos_infin);
	}
	deadline_.async_wait(boost::bind(&client::check_deadline, this));
}

std::size_t
client::receive(
	const boost::asio::mutable_buffer& buffer,
	boost::posix_time::time_duration timeout,
	boost::system::error_code& ec)
{
	deadline_.expires_from_now(timeout);

	ec = boost::asio::error::would_block;
	std::size_t length = 0;

	socket_.async_receive(boost::asio::buffer(buffer),
		boost::bind(&client::handle_receive, _1, _2, &ec, &length));

	do io_service_.run_one(); while (ec == boost::asio::error::would_block);

	return length;
}

void
client::handle_receive(
	const boost::system::error_code& ec, std::size_t length,
	boost::system::error_code* out_ec, std::size_t* out_length)
{
	*out_ec = ec; *out_length = length;
}

}
